import torch
import numpy as np
from enum import Enum

MIN_ERR = 1e-7

class OpTypes(Enum):
    NA = 0 # new standard is not available
    MOVE = 1
    RAND = 2
    CAST = 3
    COMPUTE_INTEGER = 4
    COMPUTE_QUANT = 5
    COMPUTE_FLOAT = 6
    COMPUTE_FLOAT_HIGH_PRECISION = 7
    VECTOR_FUSION = 8
    CV_FUSION = 9

def get_eb_threshold(dtype:torch.dtype):
    '''
    根据数据类型获取 eb阈值
    '''
    eb_threshold = 0
    if dtype in [torch.bfloat16]:
        eb_threshold = 2**(-7)
    if dtype in [torch.float16]:
        eb_threshold = 2**(-10)
    if dtype in [torch.float32]:
        eb_threshold = 2**(-14)
    return eb_threshold

def get_err_threshold(op_type:OpTypes, dtype:torch.dtype):
    '''
    根据算子类型+数据类型获取 eb 阈值
    '''
    err_threshold = 0
    if op_type in [OpTypes.MOVE, OpTypes.RAND, OpTypes.CAST, OpTypes.COMPUTE_INTEGER]:
        pass
    if op_type in [OpTypes.COMPUTE_QUANT, OpTypes.COMPUTE_FLOAT]:
        if dtype in [torch.bfloat16]:
            err_threshold = 2**(-7)
        if dtype in [torch.float16]:
            err_threshold = 2**(-8)
        if dtype in [torch.float32]:
            err_threshold = 2**(-11)
    if op_type in [OpTypes.CV_FUSION]:
        if dtype in [torch.bfloat16]:
            err_threshold = 2**(-8)
        if dtype in [torch.float16]:
            err_threshold = 2**(-11)
        if dtype in [torch.float32]:
            err_threshold = 2**(-14)
    return err_threshold


def ref_compare(golden:torch.Tensor, actual:torch.Tensor, err):
    '''
    单标杆、浮点比对方法|actual - expected| <= err × max(1, | expected |)
    '''
    golden = golden.to(torch.float32)
    golden_nmax = torch.clamp(torch.abs(golden), min = 1)
    abs_error = torch.abs(actual.to(torch.float32) - golden)
    result = (abs_error <= err * golden_nmax).all()
    # logging.info(f"new golden result:{result}")
    return result


def get_abs_error(golden:torch.Tensor, actual:torch.Tensor):
    '''
    绝对误差
    '''
    abs_error = torch.abs(actual.to(torch.float32) - golden.to(torch.float32))
    return abs_error


def get_rel_error(golden:torch.Tensor, actual:torch.Tensor):
    '''
    相对误差
    '''
    golden = golden.to(torch.float32)
    abs_error = get_abs_error(golden, actual)
    rel_error = abs_error / (torch.abs(golden) + MIN_ERR)
    return rel_error


def get_mare(golden:torch.Tensor, actual:torch.Tensor):
    '''
    最大相对误差：max relative error，MARE
    '''
    rel_error = get_rel_error(golden, actual)
    mare = torch.max(rel_error.flatten())
    return mare


def get_mere(golden:torch.Tensor, actual:torch.Tensor):
    '''
    平均相对误差：mean relative error，MERE
    '''
    rel_error = get_rel_error(golden, actual)
    mere = torch.mean(rel_error)
    return mere


def get_rmse(golden:torch.Tensor, actual:torch.Tensor):
    '''
    均方根误差:Root Mean Squared Error，RMSE
    '''
    golden = golden.to(torch.float32)
    sqr_err = torch.pow((actual.to(torch.float32) - golden), 2)
    rmse = torch.sqrt(torch.mean(sqr_err))
    return rmse

def get_eb(golden:torch.Tensor, actual:torch.Tensor):
    '''
    计算误差均衡性（EB）
    '''
    golden = golden.to(torch.float32)
    golden_nmax = torch.clamp(torch.abs(golden), min = 1)
    actual_error = actual.to(torch.float32) - golden
    EB = torch.mean(actual_error / golden_nmax)
    return EB


def check_operator_accuracy(
    golden: torch.Tensor,
    actual: torch.Tensor,
    mare_cpu_low: float = 0.0,
    mere_cpu_low: float = 0.0, 
    rmse_cpu_low: float = 0.0, 
    mare_rate_limit: float = 10.0, 
    mere_rate_limit: float = 2.0, 
    rmse_rate_limit: float = 2.0, 
):
    """
    检查算子输出与参考结果的精度是否满足标准
    """
    golden = golden.cpu()
    actual = actual.cpu()
    # 验证输入合法性
    assert actual.shape == golden.shape, \
        f"\Shape match failed: golden={golden.shape}, actual={actual.shape}"
    assert actual.device == golden.device, \
        f"Device match failed: golden={golden.device}, actual={actual.device}"
    assert actual.dtype == golden.dtype, \
        f"Dtype match failed: golden={golden.dtype}, actual={actual.dtype}"
    
    # 验收类型
    op_type = OpTypes.CV_FUSION
    eb_threshold = get_eb_threshold(actual.dtype)
    err_threshold = get_err_threshold(op_type, actual.dtype)

    error_threshold = err_threshold * mare_rate_limit

    # 基础误差计算
    abs_error = get_abs_error(golden, actual)
    # 总元素
    total_elements = golden.numel()
    # 计算核心误差指标
    # 相对误差
    rel_error = get_rel_error(golden, actual)
    # 最大相对误差 
    mare = get_mare(golden, actual)
    # 平均相对误差
    mere = get_mere(golden, actual)
    # 均方根误差
    rmse = get_rmse(golden, actual)
    # 误差均衡性(EB)
    eb = get_eb(golden, actual)
    # 补充：最大绝对误差、最大绝对误差
    max_abs_error = torch.max(abs_error).item()
    mean_abs_error = torch.mean(abs_error).item()
    # 相对误差率统计
    rel_error_mask = rel_error > error_threshold
    rel_error_count = torch.sum(rel_error_mask).item()
    rel_error_rate = rel_error_count / total_elements
    # 补充：绝对误差率统计
    abs_error_mask = abs_error > error_threshold
    abs_error_count = torch.sum(abs_error_mask).item()
    abs_error_rate = abs_error_count / total_elements
    # 定位最大相对误差坐标
    max_rel_indices = ()
    first_rel_coord = ()
    # if total_elements > 0:
    #     max_rel_mask = rel_error == mare
    #     max_rel_indices = torch.nonzero(max_rel_mask, as_tuple = True)
    #     first_rel_coord = tuple(coord[0].item() for coord in max_rel_indices)
    # else:
    #     first_rel_coord = ()
    if total_elements > 0:
        # mare_tensor = torch.max(rel_error)  # 保留张量形式，避免精度损失
        # mare = mare_tensor.item()  # 仅在需要数值时转换
        # max_rel_mask = rel_error == mare_tensor  # 用张量比较
        # max_rel_indices = torch.nonzero(max_rel_mask, as_tuple = True)
        # # 转换为可读的坐标格式（取第一个最大误差点）
        # first_rel_coord = tuple(coord[0].item() for coord in max_rel_indices)
        mare_tensor = torch.max(rel_error)  # 保留张量形式，避免精度损失
        mare = mare_tensor.item()  # 仅在需要数值时转换
        max_rel_mask = rel_error == mare_tensor  # 用张量比较
        max_rel_indices = torch.nonzero(max_rel_mask, as_tuple=True)
        # 检查是否有有效坐标
        if all(coord.numel() > 0 for coord in max_rel_indices):
            first_rel_coord = tuple(coord[0].item() for coord in max_rel_indices)
        else:
            first_rel_coord = ()  # 或其他默认值
    else:
        first_rel_coord = ()  
    # 定位最大绝对误差坐标
    max_abs_indices = ()
    first_abs_coord = ()
    # if total_elements > 0: 
    #     max_abs_mask = abs_error == max_abs_error
    #     max_abs_indices = torch.nonzero(max_abs_mask, as_tuple = True)
    #     first_abs_coord = tuple(coord[0].item() for coord in max_abs_indices)
    # else: 
    #     first_abs_coord = ()
    if total_elements > 0:
    # 保留张量形式的最大值，避免精度损失
        max_abs_error_tensor = torch.max(abs_error)
        max_abs_error = max_abs_error_tensor.item()  # 仅在需要数值时转换
        max_abs_mask = abs_error == max_abs_error_tensor  # 用张量比较
        max_abs_indices = torch.nonzero(max_abs_mask, as_tuple=True)
        
        # 检查是否有有效坐标，避免访问空张量
        if all(coord.numel() > 0 for coord in max_abs_indices):
            first_abs_coord = tuple(coord[0].item() for coord in max_abs_indices)
        else:
            first_abs_coord = ()  # 或其他默认值（如(-1, -1)表示无有效坐标）
    else:
        first_abs_coord = ()

    # 打印所有统计信息

    print("[Accuracy Check Result]")
    print(f"- [Tensor info] Shape = {actual.shape}, Elements num = {total_elements}, Device = {actual.device}. \n")
    
    print("[Main error metrics]")
    print(f"- [MARE(Max Reletive Error)] {mare:.8f}")
    print(f"- [MERE(Mean Reletive Error)] {mere:.8f}")
    print(f"- [RMSE(Root Mean Squared Error)] {rmse:.8f}")
    print(f"- [EB(Error Balance)] {eb:.8f}")
    # 补充：最大绝对误差、最大绝对误差
    print(f"- [Max Absolute Error] {max_abs_error:.8f}")
    print(f"- [Mean Absolute Error] {mean_abs_error:.8f}\n")
    
    print("[Error Distribution Statistics]")
    print(f"- [Error Threshold] {error_threshold:.8f}")
    print(f"- [Golden Data Range] ({torch.min(golden):.8f}, {torch.max(golden):.8f})")
    print(f"- [Actual Data Range] ({torch.min(actual):.8f}, {torch.max(actual):.8f})")
    print(f"- [Reletive Error Elements Num] {rel_error_count}")
    print(f"- [Reletive Error Rate] {rel_error_rate:.6%}")
    # 补充：绝对误差率统计
    print(f"- [Absolute Error Elements Num] {abs_error_count}")
    print(f"- [Absolute Error Rate] {abs_error_rate:.6%}\n")
    
    print("[Max Reletive Error Info]")
    if total_elements > 0: 
        is_valid = (len(first_rel_coord) == len(golden.shape) 
                    and all(0 <= coord < dim for coord, dim in zip(first_rel_coord, golden.shape)))
        if is_valid:
            print(f"- [Max Reletive Error Coord] {first_rel_coord}")
            print(f"- [Golden Data] {golden[first_rel_coord].item():.8f}")
            print(f"- [Actual Data] {actual[first_rel_coord].item():.8f}")
            print(f"- [Abs Error]{abs_error[first_rel_coord].item():.8f}")
            print(f"- [Rel Error] {rel_error[first_rel_coord].item():.8f}\n")

    print("[Max Absolute Error Info]")
    if total_elements > 0:
        print(f"- [Max Absolute Error Coord] {first_abs_coord}")
        # 检查坐标有效性：非空且维度与张量形状匹配
        is_valid = (len(first_abs_coord) == len(golden.shape) 
                    and all(0 <= coord < dim for coord, dim in zip(first_abs_coord, golden.shape)))
        if is_valid:
            print(f"- [Golden Data] {golden[first_abs_coord].item():.8f}")
            print(f"- [Actual Data] {actual[first_abs_coord].item():.8f}")
            print(f"- [Abs Error] {abs_error[first_abs_coord].item():.8f}")
            print(f"- [Rel Error] {rel_error[first_abs_coord].item():.8f}\n")
    
    
    # 验收标准检查
    mare_rate = mare / max(mare_cpu_low, err_threshold)
    mare_rate_check = mare_rate < mare_rate_limit

    mere_rate = mere / max(mere_cpu_low, err_threshold)
    mere_rate_check = mere_rate < mere_rate_limit

    rmse_rate = rmse / max(rmse_cpu_low, err_threshold)
    rmse_rate_check = rmse_rate < rmse_rate_limit

    eb_check = eb < eb_threshold
    
    print(f"[Final Pass Result] mare / max(mare_cpu_low, {err_threshold}) = {mare_rate} < {mare_rate_limit}? {mare_rate_check}")
    print(f"[Final Pass Result] mere / max(mere_cpu_low, {err_threshold}) = {mere_rate} < {mere_rate_limit}? {mere_rate_check}")
    print(f"[Final Pass Result] rmse / max(rmse_cpu_low, {err_threshold}) = {rmse_rate} < {rmse_rate_limit}? {rmse_rate_check}")
    print(f"[Final Pass Result] eb = {eb} < {eb_threshold}? {eb_check}")

    result_check = mare_rate_check and mere_rate_check and rmse_rate_check and eb_check
    # result_check = ref_compare(golden, actual, err_threshold)

    print_str = f"[{'Success' if result_check else 'Failure'}]"
    if not mare_rate_check: 
        print_str += f"mare_rate={mare_rate} < {mare_rate_limit}? {mare_rate_check}. "
    if not mere_rate_check: 
        print_str += f"mere_rate={mere_rate} < {mere_rate_limit}? {mere_rate_check}. "
    if not rmse_rate_check: 
        print_str += f"rmse_rate={rmse_rate} < {rmse_rate_limit}? {rmse_rate_check}. "
    if not eb_check: 
        print_str += f"eb={eb} < {eb_threshold}? {eb_check}"

    return {
        "result_check": result_check, 
        "print_str": print_str, 
    }


